package com.zhuangjie.mvc.ws.core;

import com.zhaungjie.base.utils.json.JsonForJSUtils;
import com.zhuangjie.mvc.ws.constant.WSConstant;
import com.zhuangjie.mvc.ws.entity.WSResult;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class WebSocketSessionStore {
    /**
     * 保存连接 session 的地方
     */
    private Map<String, Map<String, WebSocketSession>> sessionsMap = new ConcurrentHashMap<>();

    /**
     * 添加会话到sessionStore中
     * @param sessionId
     * @param session
     */
    public void put(String sessionId, WebSocketSession session) {
        String key = session.getUri().getPath();
        Map<String, WebSocketSession> sessions = sessionsMap.get(key);
        if (sessions == null) sessionsMap.put(key,sessions = new HashMap<>());
        sessions.put(sessionId,session);
    }

    /**
     * 指定集合-id删除会话
     *
     * @param key       钥匙
     * @param sessionId 会话id
     */
    public void remove(String key, String sessionId) {
        Map<String, WebSocketSession> session = sessionsMap.get(key);
        if (session == null) {
            return;
        }
        session.remove(sessionId);
    }

    /**
     * 根据session查找所在的集合
     *
     * @param session 一场
     * @return {@link Map}<{@link String},{@link WebSocketSession}>
     */
    private Map<String,WebSocketSession> findSessionRegion(WebSocketSession session){
        return this.sessionsMap.get(session.getUri().getPath());
    }

    public void remove(WebSocketSession session) {
        Map<String, WebSocketSession> sessions = findSessionRegion(session);
        if (sessions == null || sessions.isEmpty()) {
            return;
        }
        Iterator<Map.Entry<String, WebSocketSession>> iterator = sessions.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, WebSocketSession> entry = iterator.next();
            if (Objects.equals(entry.getValue(), session)) {
                iterator.remove();
            }
        }
    }


    /**
     * 接收改送消息到session的请求
     * @param key
     * @param data
     * @param sessionIds
     * @return
     * @throws IOException
     */
    public List<String> sendMessage(String key, WSResult data, String ... sessionIds) throws IOException {
        if (key == null || sessionIds == null) return null;
        // 获取工作集合（sessions）
        Map<String, WebSocketSession> sessions = sessionsMap.get(key);
        if(isNotifyAll(sessionIds)) {
            // 通知sessions里面的所有会话
            for (WebSocketSession session : sessions.values()) {
                try {
                    sendMessageCore(session, data);
                }catch (Exception e){}
            }
            return null;
        }
        // 通知sessions里面的部分
        List<String>  failed = null;
        for (String sessionId : sessionIds) {
            WebSocketSession currentSession = sessions.get(sessionId);
            if (currentSession != null) {
                sendMessageCore(currentSession,data);
            }else {
                // 如果当前节点时不存在sessionId对应的会话，则将sessionId添加到failed列表中，进行后续的广播给其它节点
                if (failed == null) failed = new ArrayList<>();
                failed.add(sessionId);
            }
        }
        return failed;
    }


    /**
     * 发送消息核心
     *
     * @param session 一场
     * @param data    数据
     * @throws IOException IOException
     */
    private void sendMessageCore(WebSocketSession session, WSResult data) throws IOException {
        session.sendMessage(new TextMessage(JsonForJSUtils.stringify(data)));
    }

    /**
     * 是否通知所有
     *
     * @param sessionIds 会话ID
     * @return boolean
     */
    private boolean isNotifyAll(String[] sessionIds) {
        return sessionIds != null && sessionIds.length > 0 && WSConstant.ANY_SESSION_ID.equals(sessionIds[0]);
    }

    /**
     * 查找session对应的serviceId(业务id)
     *
     * @param session 一场
     * @return {@link String}
     */
    public String findServiceId(WebSocketSession session) {
        if (session == null) return null;
        Map<String, WebSocketSession> sessionRegion = findSessionRegion(session);
        if (sessionRegion == null) return null;
        for (Map.Entry<String, WebSocketSession> entry : sessionRegion.entrySet()) {
            if (entry.getValue().equals(session)) {
                return entry.getKey();
            }
        }
        return null;
    }


}
